import pyaudio
import webrtcvad
import numpy as np
from pypinyin import pinyin, Style  # 如果后续需要用，可按需使用
import re
import pyttsx3
import datetime
import os
import wave
import collections

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from modelscope.pipelines import pipeline
# from deepseek_api_reply_once import get_deepseek_response # 非流式输出
from deepseek_api_connect_internet import get_deepseek_response

# 确保临时音频目录存在
os.makedirs("temp_wave", exist_ok=True)

# 参数配置
AUDIO_RATE = 16000       # 采样率（支持8000, 16000, 32000或48000）
CHUNK_SIZE = 480         # 每块大小（30ms，保证为10/20/30ms的倍数）
VAD_MODE = 1             # VAD 模式（0-3，数值越小越保守）

# 初始化 VAD
vad = webrtcvad.Vad(VAD_MODE)

# 初始化 ASR 模型
sound_recongnition_model = AutoModel(
    model=r"D:\Downloads\SenseVoiceSmall",
    trust_remote_code=False ,
    remote_code="./model.py",
    vad_model="fsmn-vad",
    vad_kwargs={"max_single_segment_time": 30000},
    # device="cuda:0",
    device="cuda:0",
    use_itn=True,
    disable_update = True,
    disable_pbar = True,
    disable_log = True
    
)

# 初始化说话人验证模型
sound_verifier_model = pipeline(
    task='speaker-verification',
    model=r'D:\Downloads\speech_campplus_sv_zh-cn_3dspeaker_16k'
)



# 把deepseek大模型输出的文字通过语音合成输出
def talkContent(content):
    """使用系统文字转语音朗读内容，提高运行效率"""
    # 初始化引擎一次，后续复用
    engine = pyttsx3.init()
    # 设置朗读速度
    engine.setProperty('rate', 160)

    # 朗读内容
    engine.say(content)
    
    # 只调用一次 runAndWait() 处理完所有加入队列的内容
    engine.runAndWait()


# 用于匹配关键词的拼音
def extract_chinese_and_convert_to_pinyin(input_string):
    """
    提取字符串中的汉字，并将其转换为拼音。
    
    :param input_string: 原始字符串
    :return: 转换后的拼音字符串
    """
    # 使用正则表达式提取所有汉字
    chinese_characters = re.findall(r'[\u4e00-\u9fa5]', input_string)
    # 将汉字列表合并为字符串
    chinese_text = ''.join(chinese_characters)
    
    # 转换为拼音
    pinyin_result = pinyin(chinese_text, style=Style.NORMAL)
    # 将拼音列表拼接为字符串
    pinyin_text = ' '.join([item[0] for item in pinyin_result])
    
    return pinyin_text

def calibrate(stream, calibration_seconds=2, chunk_duration_ms=30):
    """
    校准背景噪音：录制指定时长的音频，计算平均幅值与标准差，从而设置自适应阈值
    参数:
      calibration_seconds: 校准时间（秒）
      chunk_duration_ms: 每块时长（毫秒）
    返回:
      amplitude_threshold: 设定的音频幅值阈值
    """
    print("开始校准背景噪音，请保持安静...")
    amplitudes = []
    num_frames = int(calibration_seconds * (1000 / chunk_duration_ms))
    for _ in range(num_frames):
        audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        amplitudes.append(np.abs(audio_data).mean())
    mean_noise = np.mean(amplitudes)
    std_noise = np.std(amplitudes)
    amplitude_threshold = mean_noise + 2 * std_noise
    print(f"校准完成：噪音均值={mean_noise:.2f}，标准差={std_noise:.2f}，设置阈值={amplitude_threshold:.2f}")
    return amplitude_threshold

# 保存音频数据为WAV文件 
def save_audio_to_wav(audio_data, sample_rate, channels, filename):
    """
    将音频数据保存为WAV文件
    
    参数:
        audio_data: 字节数组格式的音频数据
        sample_rate: 采样率
        channels: 通道数
        filename: 要保存的文件名
    """
    if not audio_data:
        return
        
    # 确保目录存在
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    
    # 直接保存
    with wave.open(filename, 'wb') as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(2)  # 16位采样宽度 2*8=16
        wf.setframerate(sample_rate)
        wf.writeframes(audio_data)
    
    # print(f"已保存音频文件: {filename}")

class SpeechDetector:
    """
    SpeechDetector 负责处理音频块，结合能量预处理、VAD 和频谱分析进行语音检测，
    并在检测到语音结束后调用 ASR 模型进行转写，返回识别结果文本。
    """
    def __init__(self, amplitude_threshold):
        self.amplitude_threshold = amplitude_threshold

        # 音频缓冲区，用于存储所有音频数据，包括非语音部分
        self.speech_buffer = bytearray()
        
        # 使用循环缓冲区存储最近的所有音频数据
        self.audio_buffer = collections.deque(maxlen=500)  # 存储约15秒的音频

        # 连续帧状态，用于平滑判断语音是否开始/结束
        self.speech_state = False          # True：正在录入语音；False：非语音状态
        self.consecutive_speech = 0        # 连续语音帧计数
        self.consecutive_silence = 0       # 连续静音帧计数
        self.required_speech_frames = 2    # 连续语音帧达到此值后确认进入语音状态（例如 2 帧大约 60ms）
        self.required_silence_frames = 34*1  # 连续静音帧达到此值后确认语音结束（例如 15 帧大约 450ms）
        self.long_silence_frames = 34*5    # 连续静音帧达到此值后确认语音结束（例如 34 帧大约 1s）

    def analyze_spectrum(self, audio_chunk):
        """
        通过频谱分析检测语音特性：
          1. 对音频块应用汉宁窗后计算 FFT
          2. 统计局部峰值数量（峰值必须超过均值的1.5倍）
          3. 当峰值数量大于等于3时，认为该块具有语音特征
        """
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        if len(audio_data) == 0:
            return False

        # 应用汉宁窗减少 FFT 泄露
        window = np.hanning(len(audio_data))
        windowed_data = audio_data * window

        # 计算 FFT 并取正频率部分
        spectrum = np.abs(np.fft.rfft(windowed_data))
        spectral_mean = np.mean(spectrum)

        peak_count = 0
        for i in range(1, len(spectrum) - 1):
            if (spectrum[i] > spectrum[i - 1] and 
                spectrum[i] > spectrum[i + 1] and 
                spectrum[i] > spectral_mean * 1.5):
                peak_count += 1

        return peak_count >= 3

    def is_speech(self, audio_chunk):
        """
        判断当前音频块是否包含语音：
          1. 先通过能量阈值预过滤低幅值数据
          2. 再结合 VAD 检测与频谱分析判断
        """
        threshold = self.amplitude_threshold if self.amplitude_threshold is not None else 11540.82
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        amplitude = np.abs(audio_data).mean()
        if amplitude < threshold:
            return False

        # 降低VAD判断条件的严格性，只检查VAD结果或频谱结果之一
        vad_result = vad.is_speech(audio_chunk, AUDIO_RATE)
        spectral_result = self.analyze_spectrum(audio_chunk)

        return vad_result and spectral_result  # 使用或操作而不是与操作，降低检测条件

    def process_chunk(self, audio_chunk):
        """
        处理每个音频块，并在识别到语音结束后返回文本结果。
        修改后的工作流程：
          - 每个音频块都会被添加到循环缓冲区中
          - 对当前块进行语音检测
          - 根据语音状态管理并处理录音数据
          - 返回识别结果和完整的音频数据
        """
        recognized_text = None
        temp_speech_buffer = None
        
        # 始终将音频块添加到循环缓冲区，确保所有音频数据都被保存
        self.audio_buffer.append(audio_chunk)
        
        is_speech_chunk = self.is_speech(audio_chunk)

        if is_speech_chunk:
            self.consecutive_speech += 1
            self.consecutive_silence = 0
            if not self.speech_state and self.consecutive_speech >= self.required_speech_frames:
                # 检测到语音开始，初始化speech_buffer
                self.speech_state = True
                self.speech_buffer = bytearray()
                
                # 将最近缓冲区中的数据（包括语音前部分）加入到speech_buffer
                for chunk in list(self.audio_buffer)[-10:]:  # 添加前300ms的音频
                    self.speech_buffer.extend(chunk)
                
                print("检测到语音开始")
            
            # 在语音状态下，继续添加音频块
            if self.speech_state:
                self.speech_buffer.extend(audio_chunk)

        else:  # 当前块不是语音
            self.consecutive_silence += 1
            self.consecutive_speech = 0
            
            # 在语音状态下，继续添加非语音块（保证完整性）
            if self.speech_state:
                self.speech_buffer.extend(audio_chunk)
            
            # 长时间无语音处理
            if self.consecutive_silence >= self.long_silence_frames and not self.speech_state:
                recognized_text = "长时间无语音"
                self.consecutive_silence = 0
                return recognized_text, None
    
            # 语音结束判断
            if self.speech_state and self.consecutive_silence >= self.required_silence_frames:
                if len(self.speech_buffer) > CHUNK_SIZE * 5:
                    print(f"\n======采集到的语音数据长度：{len(self.speech_buffer)}")

                    # 临时保存当前语音段
                    temp_speech_buffer = bytes(self.speech_buffer)

                    # 保存音频数据到WAV文件
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    wav_filename = f"temp_wave/waving_{timestamp}.wav"
                    
                    try:
                        save_audio_to_wav(temp_speech_buffer, AUDIO_RATE, 1, wav_filename)
                        recognized_text = self.sound2text(wav_filename)
                        os.remove(wav_filename)  # 使用os.remove删除单个文件

                    except Exception as e:
                        print(f"保存音频文件时出错或者删除临时文件时出错: {e}")

                # 重置状态，准备下一段语音录入
                self.speech_state = False
                self.speech_buffer = bytearray()

        return recognized_text, temp_speech_buffer


    def sound2text(self, audio_data):
        """
        利用语音识别模型将音频数据转换为文本
        """
        res = sound_recongnition_model.generate(
            input=audio_data,
            cache={},
            language="zh",        # 支持 "zh", "en", 等
            use_itn=True,
            batch_size_s=60,
            merge_vad=True,
            merge_length_s=15,
        )
        text_string = rich_transcription_postprocess(res[0]["text"])
        return text_string 

def ASR_API(stream, detector):
    """
    ASR_API 函数：打开音频流，校准背景噪音，并持续监听语音。
    每次处理音频块后，如果识别到结果，则实时输出文本识别结果。
    """
    while True:
        # 从流中准确读取一个音频块
        audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
        
        # 处理音频块
        chinese_string, wave_data = detector.process_chunk(audio_chunk)
        
        # 若识别到结果，则输出（无结果时不输出）
        if chinese_string:
            yield chinese_string, wave_data


if __name__ == "__main__":
    import time
    p = pyaudio.PyAudio()
    stream = p.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=AUDIO_RATE,
        input=True,
        frames_per_buffer=CHUNK_SIZE
    )

    # 校准背景噪音，获取幅值阈值
    amplitude_threshold = calibrate(stream) # 11540.82
    detector = SpeechDetector(amplitude_threshold=amplitude_threshold)
    key_word_pinyin = "ni hao xiao pang"
    print("开始监听，请开始说话...(按 Ctrl+C 停止)")


    while True:
        print("=====================开始新的一轮监听=====================")


        #============================================================================================================#
        # 验证说话人的关键词和声纹
        for chinese_string, audio_byte_data  in ASR_API(stream, detector):
            print("识别结果:", chinese_string)
            # print("音频数据:", len(wave_numpy_data))

            chinese_pinyin = extract_chinese_and_convert_to_pinyin(chinese_string)
            if key_word_pinyin in chinese_pinyin:

                print("================================识别到关键词，开始识别===================================")

                # 调用模型进行声纹识别
                recongnize_result = sound_verifier_model([r"D:\Project\Chat_Project\my_recording.wav", np.frombuffer(audio_byte_data, dtype=np.int16)])
                print("声纹识别结果:", recongnize_result)
                if recongnize_result["score"] > 0.4:
                    print("===================================声纹识别成功，开始对话===================================")
                    
                    # 关键词和声纹验证成功，开始对话，进入下一阶段
                    talkContent("你好，请开始提问吧！")
                    break
                else:
                    talkContent("声纹识别失败，请稍后再试！")
        #============================================================================================================#


        #============================================================================================================#
        # 说话人验证成功后，用于接受语音输入并通过大模型输出结果，2秒无输入则跳过该循环重新进入说话人验证阶段
        for recognized_text, _ in ASR_API(stream, detector):
            print("识别结果:", recognized_text)

            if recognized_text == "长时间无语音": # 长时间无语音
                talkContent("长时间无语音，即将退出智能对话！")
                print("===============识别结果:", recognized_text, "中断识别!")
                break

            else: #把识别到的结果传给DS大模型，让DS大模型输出结果
                deepseek_response_text = get_deepseek_response(recognized_text)
                print("DS大模型输出结果:", deepseek_response_text)
                talkContent(deepseek_response_text)
                talkContent("我的回答完毕，继续提问或5秒后自动退出智能对话！")
            
            print("\n")
        print("\n")
        #============================================================================================================#


        # 打断功能的设计
        # 针对打断的功能，可以考虑通过检测声音的阈值，如果检测超过1秒的声音，则中断播放。
        # 或者检测关键词“暂停”、“停止”等
        # 在打断时需要考虑麦克风的设备的占用问题

    # 回收资源，当需要把音频流资源分配给其他的程序时，需要调用“stream.stop_stream()和stream.close()”
    stream.stop_stream()
    stream.close()
    p.terminate()
